{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f854b3ac-ffa2-47f6-9b6e-1eb9a029af22",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import datasets\n",
    "from dataclasses import dataclass\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import torch\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cb5e15e-bf27-4202-b18e-cb292fc6e307",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda'\n",
    "model_path = './Qwen2.5-0.5B'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f436ff45-fc6e-482f-9f77-10fecd436ef8",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_path,\n",
    "    dtype='auto',\n",
    "    device_map='auto'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a551290e-0a32-4448-96c2-976318b0957b",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b28b0de2-e4b8-4f1a-a2dd-a37e3488357a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.generation_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5afc74c7-dc24-4dd9-bdeb-a3cf11bb21d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.generation_config.do_sample = True\n",
    "model.generation_config.eos_token_id = [151645, 151643]\n",
    "model.generation_config.pad_token_id = 151643\n",
    "model.generation_config.temperature = 0.7\n",
    "model.generation_config.top_p = 0.8\n",
    "model.generation_config.top_k = 20\n",
    "model.generation_config.repetition_penalty = 1.05"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2d844dc-a69a-4342-930f-980744f54097",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.generation_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "967cc8b5-14d7-4862-8e52-458b9b43754a",
   "metadata": {},
   "outputs": [],
   "source": [
    "chat = [\n",
    "    {\"role\": \"user\", \"content\": \"《滕王阁序》是谁写的？\"}\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dc6e818-4467-4569-93ec-b1bb2fa74bf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import pipeline\n",
    "\n",
    "pipeline = pipeline(task=\"text-generation\", model=model_path, dtype=\"auto\", device_map=\"auto\")\n",
    "response = pipeline(chat, max_new_tokens=512)\n",
    "print(response[0][\"generated_text\"][-1][\"content\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91de7fd7-5270-4924-8aad-91bdfcad879a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
